import datasets
from transformers import AutoTokenizer, AutoModelForSequenceClassification, get_linear_schedule_with_warmup

train_dataset, test_dataset = datasets.load_dataset('imdb', split=['train', 'test'])
model = AutoModelForSequenceClassification.from_pretrained('/root/.cache/huggingface/transformers', num_labels=2)
tokenizer = AutoTokenizer.from_pretrained('bert-base-cased')

